package com.bleeth.flow.step.row;

import cn.hutool.core.collection.CollUtil;
import com.bleeth.flow.core.common.APlugin;
import com.bleeth.flow.core.common.PluginAnnotation;
import com.bleeth.flow.core.common.PluginTypeEnum;
import com.jd.platform.async.worker.WorkResult;
import com.jd.platform.async.wrapper.WorkerWrapper;
import lombok.Data;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.functions;

import java.io.Serializable;
import java.util.List;
import java.util.Map;

/**
 * @author ：Bleeth
 * @date ：2021-08-06
 * @description：数据合并组件
 */
@Data
@PluginAnnotation(name = "数据合并组件",
        type = PluginTypeEnum.ROW,
        description = "",
        id = "UnionPlugin")
public class UnionPlugin extends APlugin implements Serializable {

    private static final long serialVersionUID = -8611579221464491033L;


    private List<UnionParamBean> unionParamList;



    @Override
    public Dataset<Row> action(Dataset<Row> param, Map<String, WorkerWrapper> allWrappers) {
        String leftPluginName = fromList.get(0);
        String rightPluginName = fromList.get(1);

        WorkerWrapper leftWrapper = allWrappers.get(leftPluginName);
        WorkerWrapper rightWrapper = allWrappers.get(rightPluginName);

        WorkResult leftWorkResult = leftWrapper.getWorkResult();
        WorkResult rightWorkResult = rightWrapper.getWorkResult();

        Dataset<Row> leftDs = null;
        Dataset<Row> rightDs = null;

        leftDs = (Dataset<Row>) leftWorkResult.getResult();
        rightDs = (Dataset<Row>) rightWorkResult.getResult();

        String[] leftColumns = leftDs.columns();
        String[] rightColumns = rightDs.columns();

        List<String> leftColList = CollUtil.newArrayList(leftColumns);
        List<String> rightColList = CollUtil.newArrayList(rightColumns);

        for (int i = 0; i < unionParamList.size(); i++) {
            UnionParamBean unionParamBean = unionParamList.get(i);
            String leftFieldName = unionParamBean.getLeftFieldName();
            String rightFieldName = unionParamBean.getRightFieldName();
            String unionName = unionParamBean.getUnionName();
            if (leftColList.contains(leftFieldName)) {
                leftDs = leftDs.withColumnRenamed(leftFieldName, unionName);
            }
            if (rightColList.contains(rightFieldName)) {
                rightDs = rightDs.withColumnRenamed(rightFieldName, unionName);
            }
        }

        leftColList = CollUtil.newArrayList(leftDs.columns());
        rightColList = CollUtil.newArrayList(rightDs.columns());
        List<String> unionList = CollUtil.distinct(CollUtil.addAll(leftColList, rightColList));
        leftColList = CollUtil.newArrayList(leftDs.columns());
        for (int i = 0; i < unionList.size(); i++) {
            String fieldName = unionList.get(i);
            if (!leftColList.contains(fieldName)) {
                leftDs = leftDs.withColumn(fieldName, functions.lit(""));
            }
            if (!rightColList.contains(fieldName)) {
                rightDs = rightDs.withColumn(fieldName, functions.lit(""));
            }
        }


        leftDs.show();

        //合并名称
        Dataset<Row> ds = leftDs.unionByName(rightDs);
        return ds;
    }


    @Data
    public static class UnionParamBean implements Serializable {
        private static final long serialVersionUID = -5117687853827849740L;

        //左字段名称
        private String leftFieldName;

        //右字段名称
        private String rightFieldName;

        //合并后的新名称
        private String unionName;
    }


}
